# 抽取特征到文件
# pickle: featurelen x 201

from Config.Config import DPN107_RGB_200_PATH

import pickle
import numpy as np
import h5py
import os

import torch
import torch.nn as nn
from torch.autograd import Variable

from FeatureExtruct.DPN107_RGB.RGB_SinglePicture_Dataset import RGB_Single_Frame_Dataset

# laod model

model = torch.load('/mnt/md1/Experiments/DPN_Extruct_200_Test1/raw_dpn107_model.pkl')
model = nn.DataParallel(model).cuda()
softmax = nn.Softmax(dim=-1).cuda()
model.load_state_dict(torch.load('/mnt/md1/Experiments/DPN_Extruct_200_Test1/backup/dpn107_RGB_fineturn_201_model_011.ckpt'))

# open eval model
model = model.eval()
softmax = softmax.eval()

# extruct feature
# def ExtructFeature(gen):
#
#     lastvid = None
#     store_preds = []
#     store_labels = []
#
#     for i,[vid,images,labels] in enumerate(gen):
#
#         if lastvid is None: lastvid = vid
#
#         feature_path = os.path.join(DPN107_RGB_200_PATH,'{}.h5'.format(vid))
#         if os.path.exists(feature_path): continue
#
#         images = Variable(images,volatile=True).cuda()
#         pred = softmax(model(images))
#
#         store_preds.append(pred.data.cpu().numpy())
#         store_labels.extend(labels)
#
#         if vid!=lastvid:
#             store_preds = np.concatenate(store_preds)
#             # save feature
#             with h5py.File(feature_path,'w') as f:
#                 f['feature201'] = store_preds
#                 f['label'] = store_labels
#                 f.attrs['vid'] = vid
#
#             lastvid = vid
#             store_preds = []
#             store_labels = []
#
#     if len(store_labels) != 0 :
#         store_preds = np.concatenate(store_preds)
#         with h5py.File(feature_path,'w') as f:
#             f['feature201'] = store_preds
#             f['labels'] = store_labels
#             f.attrs['vid'] = vid

def ExtructFeature(dataset):

    vids = dataset.vids

    for idx,vid in enumerate(vids):

        feature_path = os.path.join(DPN107_RGB_200_PATH,'{}.h5'.format(vid))
        if os.path.exists(feature_path) == True:
            continue

        if os.path.exists('/mnt/md1/Dataset/ActivityNet/Frames/{}'.format(vid)) == False:
            continue

        try:
            gen = dataset.enum_vid_feature(vid)
        except Exception:
            continue

        store_preds = []
        store_labels = []

        for i,[_,images,labels] in enumerate(gen):
            images = Variable(images,volatile=True).cuda()
            pred = softmax(model(images))
            store_preds.append(pred.data.cpu().numpy())
            store_labels.extend(labels)

        if len(store_preds) == 0: continue
        store_preds = np.concatenate(store_preds)

        with h5py.File(feature_path,'w') as f:
            f['feature201'] = store_preds
            f['label'] = store_labels
            f.attrs['vid'] = vid

        if idx%10==0: print(idx,'...')

# load dataset
# dataset = RGB_Single_Frame_Dataset('training')
# ExtructFeature(dataset)

dataset = RGB_Single_Frame_Dataset('validation')
ExtructFeature(dataset)
